from typing import Dict, Tuple, Optional
import torch
import torch.nn as nn
import torch.optim as optim

class BaseTrainer:
    def __init__(self, model: nn.Module, device: str, params: Dict):
        self.model = model.to(device)
        self.device = device
        self.params = params
        
        # 设置优化器
        self.optimizer = self._get_optimizer()
    
    def _get_optimizer(self) -> torch.optim.Optimizer:
        """根据参数配置获取优化器"""
        optimizer_name = self.params.get("optimizer", "adam").lower()
        lr = self.params.get("learning_rate", 0.001)
        
        if optimizer_name == "adam":
            return optim.Adam(self.model.parameters(), lr=lr)
        elif optimizer_name == "sgd":
            return optim.SGD(self.model.parameters(), lr=lr)
        elif optimizer_name == "rmsprop":
            return optim.RMSprop(self.model.parameters(), lr=lr)
        else:
            raise ValueError(f"Unsupported optimizer: {optimizer_name}")
    
    def train_step(self, x: torch.Tensor, y: torch.Tensor) -> Tuple[float, Optional[float]]:
        """训练一个批次，返回损失和准确率"""
        raise NotImplementedError
    
    def evaluate(self, x: torch.Tensor, y: torch.Tensor) -> Tuple[float, Optional[float]]:
        """评估一个批次，返回损失和准确率"""
        raise NotImplementedError 